
#include "Preprocess_CMP_RT.h"
#include "utool_AudioRT/Audio_RT_Process.h"
#include "utool_AudioRT/Config.h"

// 实现Mel滤波器组创建函数

std::vector<std::vector<float>> create_mel_filter_bank(int sample_rate, int n_fft, int n_mels) {
    if (sample_rate <= 0) throw std::invalid_argument("sample_rate must be > 0");
    if (n_fft <= 0) throw std::invalid_argument("n_fft must be > 0");
    if (n_mels <= 0) throw std::invalid_argument("n_mels must be > 0");

    const double fmin = 0.0;
    const double fmax = static_cast<double>(sample_rate) / 2.0;
    const int n_freqs = n_fft / 2 + 1;

    // HTK mel <-> hz
    auto hz_to_mel_htk = [](double f) {
        return 2595.0 * std::log10(1.0 + f / 700.0);
    };
    auto mel_to_hz_htk = [](double m) {
        return 700.0 * (std::pow(10.0, m / 2595.0) - 1.0);
    };

    // mel points (n_mels + 2)
    double mel_min = hz_to_mel_htk(fmin);
    double mel_max = hz_to_mel_htk(fmax);
    std::vector<double> mels(n_mels + 2);
    for (int i = 0; i < (int)mels.size(); ++i) {
        double frac = static_cast<double>(i) / static_cast<double>(n_mels + 1);
        mels[i] = mel_min + frac * (mel_max - mel_min);
    }

    // mel -> hz
    std::vector<double> hz_points(mels.size());
    for (size_t i = 0; i < mels.size(); ++i) hz_points[i] = mel_to_hz_htk(mels[i]);

    // precompute fft bin frequencies (Hz): freq[k] = k * sr / n_fft, k=0..n_freqs-1
    std::vector<double> fft_freqs(n_freqs);
    for (int k = 0; k < n_freqs; ++k) {
        fft_freqs[k] = static_cast<double>(k) * static_cast<double>(sample_rate) / static_cast<double>(n_fft);
    }

    // prepare mel bank
    std::vector<std::vector<float>> mel_fb(n_mels, std::vector<float>(n_freqs, 0.0f));

    // build each mel filter by interpolating on freq (Hz)
    for (int m = 0; m < n_mels; ++m) {
        double left_hz = hz_points[m];
        double center_hz = hz_points[m + 1];
        double right_hz = hz_points[m + 2];

        // rising side: left_hz .. center_hz (inclusive on center)
        if (center_hz > left_hz) {
            for (int k = 0; k < n_freqs; ++k) {
                double f = fft_freqs[k];
                if (f >= left_hz && f <= center_hz) {
                    double val = (f - left_hz) / (center_hz - left_hz);
                    if (val < 0.0) val = 0.0;
                    if (val > 1.0) val = 1.0;
                    mel_fb[m][k] = static_cast<float>(val);
                }
            }
        }
        // falling side: center_hz .. right_hz (exclusive on center, inclusive on right)
        if (right_hz > center_hz) {
            for (int k = 0; k < n_freqs; ++k) {
                double f = fft_freqs[k];
                if (f > center_hz && f <= right_hz) {
                    double val = (right_hz - f) / (right_hz - center_hz);
                    if (val < 0.0) val = 0.0;
                    if (val > 1.0) val = 1.0;
                    mel_fb[m][k] = static_cast<float>(val);
                }
            }
        }
    }

    // Slaney-style normalization (enorm = 2 / (hz[m+2] - hz[m]))
    for (int m = 0; m < n_mels; ++m) {
        double hz_diff = hz_points[m + 2] - hz_points[m];
        if (hz_diff <= 0.0) continue;
        double enorm = 2.0 / hz_diff;
        for (int k = 0; k < n_freqs; ++k) {
            mel_fb[m][k] = static_cast<float>(mel_fb[m][k] * enorm);
        }
    }

    return mel_fb;
}


// 主函数：计算Mel谱
void compute_Melstft(std::vector<float> &mono, 
                     std::vector<std::vector<float>> &mel_spectrogram,
                     std::vector<std::vector<float>> &mel_filter_bank) 
{
    // 参数设置（与Python代码一致）
    const int n_fft = WINDOW_SIZE_ASC;
    const int hop_length = HOP_SIZE_ASC;
    const int win_length = WINDOW_SIZE_ASC;
    const int sample_rate = SAMPLE_RATE;
    const int n_mels = NUM_MELS;
    const float epsilon = 1e-5f;  // 与Python代码一致

    // 4. 生成分析窗口
    std::vector<__fp16> Anly_Windows(win_length);
    Anly_Windows = asymmetric_Analy_windows(hop_length,n_fft,10);

    // 5. 信号填充（中心对齐）
    std::vector<float> padded(mono.size() + n_fft, 0.0f);
    std::copy(mono.begin(), mono.end(), padded.begin() + n_fft/2);

    // 6. 分帧处理
    std::vector<std::vector<float>> frames;
    for (size_t start = 0; start + win_length <= padded.size(); start += hop_length) {
        std::vector<float> frame(win_length);
        for (int i = 0; i < win_length; ++i) {
            frame[i] = padded[start + i] * Anly_Windows[i];
        }
        frames.emplace_back(std::move(frame));
    }

    // 7. 创建Mel滤波器组
    // std::vector<std::vector<float>> mel_filter_bank = create_mel_filter_bank(sample_rate, n_fft, n_mels);

    // 8. 初始化结果：mel_spectrogram，大小为n_frames x n_mels
    mel_spectrogram.resize(frames.size(), std::vector<float>(n_mels, 0.0f));
    // std::vector<std::vector<float>> mel_spectrogram(frames.size(), std::vector<float>(n_mels, 0.0f));

    // 9. FFT配置
    kiss_fft_cfg cfg = kiss_fft_alloc(n_fft, 0, NULL, NULL);
    if (!cfg) {
        throw std::runtime_error("Failed to initialize FFT");
    }

    // 10. 处理每帧
    for (size_t i = 0; i < frames.size(); ++i) {
        kiss_fft_cpx in[n_fft], out[n_fft];
        // 填充FFT输入
        for (int j = 0; j < n_fft; ++j) {
            in[j].r = j < win_length ? frames[i][j] : 0.0f;
            in[j].i = 0.0f;
        }

        // 执行FFT
        kiss_fft(cfg, in, out);

        // 计算功率谱（只取前n_fft/2+1个点）
        std::vector<float> power_spectrum(n_fft/2+1);
        for (int j = 0; j <= n_fft/2; ++j) {
            float mag = std::hypot(out[j].r, out[j].i);
            power_spectrum[j] = mag * mag;
        }

        // 应用Mel滤波器组
        for (int m = 0; m < n_mels; ++m) {
            float energy = 0.0f;
            for (int j = 0; j <= n_fft/2; ++j) {
                energy += mel_filter_bank[m][j] * power_spectrum[j];
            }
            // 对数变换（与Python代码一致）
            mel_spectrogram[i][m] = std::log(energy + epsilon);
        }
    }

    // 11. 清理资源
    kiss_fft_free(cfg);

    // return mel_spectrogram;
}

